#ifndef EYE_N_D_H
#define EYE_N_D_H

#include "kernel_operator.h"

namespace Eye {
using namespace AscendC;

constexpr int32_t BYTE_BLOCK = 32;

constexpr float POS_ONE = 1.0;

template <typename T>
class EyeND {
public:
    __aicore__ inline EyeND() = default;
    __aicore__ inline void Init(GM_ADDR y, const EyeTilingData* __restrict tilingData);
    __aicore__ inline void ProcessSmallM();
    __aicore__ inline void ProcessMidM();

private:
    template <typename T1, typename T2>
    __aicore__ inline T1 CeilDiv(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : (a + bTemp - 1) / bTemp;
    };

    template <typename T1, typename T2>
    __aicore__ inline T1 CeilAlignA2B(T1 a, T2 b) {
        T1 bTemp(b);
        return bTemp == 0 ? a : CeilDiv(a, bTemp) * bTemp;
    };

    __aicore__ inline void CopyOut(int64_t offset, int64_t dataCount);
private:
    TPipe pipe;
    // TQue<QuePosition::VECIN, 1> xQue;
    TQue<QuePosition::VECOUT, 1> yQue;
    GlobalTensor<T> xGM, yGM;
    TBuf<QuePosition::VECCALC> tempValBuf;
    LocalTensor<float> tempValLT;
    int64_t blockIdx = 0;
    uint64_t perBlockCount = 0;

    // tiling params
    uint64_t numRows = 0;
    uint64_t numColumns = 0;
    uint64_t batchNum = 0;
    uint64_t taskPerCore = 0;
    uint64_t taskTail = 0;
    uint64_t tilingMode = 0;
    uint64_t tilingCoreNum = 0;

    uint64_t numOnesPerBatch = 0;
    uint64_t offsetOne = 0;
    uint64_t numElemPerBatch = 0;
    uint64_t ubMaxProcCount = 0;
    uint64_t numElemPerBlock = 0;
};

template <typename T>
__aicore__ inline void EyeND<T>::Init(GM_ADDR y,
                                                        const EyeTilingData* __restrict tilingData) {
    blockIdx = GetBlockIdx();
    perBlockCount = BYTE_BLOCK / sizeof(T);

    numRows = tilingData->numRows;
    numColumns = tilingData->numColumns;
    batchNum = tilingData->batchNum;
    taskPerCore = tilingData->taskPerCore;
    taskTail = tilingData->taskTail;
    tilingMode = tilingData->tilingMode;
    tilingCoreNum = tilingData->tilingCoreNum;
    ubMaxProcCount = tilingData->ubMaxProcCount;
    numElemPerBlock = tilingData->numElemPerBlock;

    numOnesPerBatch = numColumns;
    if (numRows < numColumns) {
        numOnesPerBatch = numRows;
    }

    // the base offset of each one in output_y_gm
    offsetOne = numColumns + 1;
    numElemPerBatch = numRows * numColumns;

    // xGM.SetGlobalBuffer((__gm__ T*)x, numElemPerBatch);
    yGM.SetGlobalBuffer((__gm__ T*)y, numElemPerBatch);

    uint64_t singleBufferSize = ubMaxProcCount;
    
    pipe.InitBuffer(yQue, 1, singleBufferSize);

    // printf(
    //     "==zf==Kernel Tiling Data, ubMaxProcCount:%lu, totalDataCount:%lu, binsValue:%ld, minValue:%f, maxValue:%f, "
    //     "isMinMaxEqual:%ld, loopTime:%lu, tailCount:%lu\n ",
    //     ubMaxProcCount, totalDataCount, binsValue, minValue, maxValue, isMinMaxEqual, loopTime, tailCount);
}

template <typename T>
__aicore__ inline void EyeND<T>::ProcessSmallM() { 
    uint64_t alignCount = CeilAlignA2B(numElemPerBatch, perBlockCount);
    LocalTensor<T> yLT = yQue.AllocTensor<T>();
    Duplicate(yLT, (T) 0.0, alignCount);
    uint64_t offset = 0;
    for (int i=0; i < numOnesPerBatch; i++) {
        offset = i * (numColumns + 1);
        yLT.SetValue(offset, (T) 1.0);
    }
    yQue.EnQue(yLT);
    CopyOut(0, alignCount);
}

template <typename T>
__aicore__ inline void EyeND<T>::ProcessMidM() { 
    
    // 32B aligned number of elements, max size = 32B * 32B
        // aligned_num_elem = (self.ceil_div(self.num_ones_per_batch * self.num_columns, self.num_elem_per_block) * 
        //                     self.num_elem_per_block)
    offsetLastOne = (numOnesPerBatch - 1) * (numColumns + 1)
    uint64_t alignedNumElem = CeilDiv(numOnesPerBatch * numColumns, numElemPerBlock) * numElemPerBlock;
    uint64_t numElem = numOnesPerBatch * numColumns;
    uint64_t frontNumOnes = (numElem - 1) / (numColumns + 1) + 1;
    uint64_t offsetGm = 0;
    uint64_t tailNumOnes = 0;
    LocalTensor<T> yLT = yQue.AllocTensor<T>();
    Duplicate(yLT, (T) 0.0, alignedNumElem);
    uint64_t offset = 0;
    for (int i=0; i < frontNumOnes; i++) {
        offset = i * (numColumns + 1);
        yLT.SetValue(offset, (T) 1.0);
    }
    yQue.EnQue(yLT);
    CopyOut(0, alignedNumElem);
    offsetGm = numElem - numElemPerBlock;
    tailNumOnes = numOnesPerBatch - frontNumOnes;
    if (tailNumOnes > 0) {
        
        LocalTensor<T> yLT_ = yQue.AllocTensor<T>();
        Duplicate(yLT_, (T) 0.0, numElemPerBlock);
        for (int i=0; i < tailNumOnes; i++) {
            offset = (frontNumOnes + i) * (numColumns + 1) - offsetGm;
            yLT_.SetValue(offset, (T) 1.0);
        }
        
        yQue.EnQue(yLT_);
        CopyOut(offsetGm, numElemPerBlock);
    }
}

template <typename T>
__aicore__ inline void EyeND<T>::ProcessLargeM() {     
    // 32B aligned number of elements, max size = 32B * 32B
        // aligned_num_elem = (self.ceil_div(self.num_ones_per_batch * self.num_columns, self.num_elem_per_block) * 
        //                     self.num_elem_per_block)
    uint64_t alignedNumElem = CeilDiv(numOnesPerBatch * numColumns, numElemPerBlock) * numElemPerBlock;
    uint64_t numElem = numOnesPerBatch * numColumns;
    uint64_t frontNumOnes = (numElem - 1) / (numColumns + 1) + 1;
    uint64_t offsetGm = 0;
    uint64_t tailNumOnes = 0;
    LocalTensor<T> yLT = yQue.AllocTensor<T>();
    Duplicate(yLT, (T) 0.0, alignedNumElem);
    uint64_t offset = 0;
    for (int i=0; i < frontNumOnes; i++) {
        offset = i * (numColumns + 1);
        yLT.SetValue(offset, (T) 1.0);
    }
    yQue.EnQue(yLT);
    CopyOut(0, alignedNumElem);
    offsetGm = numElem - numElemPerBlock;
    tailNumOnes = numOnesPerBatch - frontNumOnes;
    if (tailNumOnes > 0) {
        
        LocalTensor<T> yLT_ = yQue.AllocTensor<T>();
        Duplicate(yLT_, (T) 0.0, numElemPerBlock);
        for (int i=0; i < tailNumOnes; i++) {
            offset = (frontNumOnes + i) * (numColumns + 1) - offsetGm;
            yLT_.SetValue(offset, (T) 1.0);
        }
        
        yQue.EnQue(yLT_);
        CopyOut(offsetGm, numElemPerBlock);
    }
}

template <typename T>
__aicore__ inline void EyeND<T>::CopyOut(int64_t offset, int64_t dataCount) {
    
    LocalTensor<T> yOutLT = yQue.DeQue<T>();
    uint64_t offsetBatch = 0;
    for (int j=0; j < batchNum; j++) {
        
        offsetBatch = j * numElemPerBatch;
        DataCopy(yGM[offsetBatch + offset], yOutLT, dataCount);
    
    }
    yQue.FreeTensor(yOutLT);
}
}  // namespace Eye

#endif  // EYE_N_D_H